package sf.codegen;

import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.JavassistUtils;
import javassist.LoaderClassPath;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.smallframework.spring.path.ClassPathScanningUtils;
import sf.tools.IOUtils;
import sf.tools.SystemUtils;

import java.io.File;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * Entity静态增强任务类,按URL来访问资源
 * <h3>作用</h3> 这个类中提供了{@link #enhance(String...)}
 * 方法，可以对当前classpath下的Entity类进行字节码增强。
 * @author
 * @Date
 */
public class EntityEnhancerJavassist implements IEntityEnhancer {
    private static final Logger logger = LoggerFactory.getLogger(EntityEnhancerJavassist.class);

    private List<URL> roots;

    @Override
    public IEntityEnhancer addRoot(URL url) {
        if (url != null) {
            if (roots == null) {
                roots = new ArrayList<>();
            }
            roots.add(url);
        }
        return this;
    }

    /**
     * 在当前的classpath目录下扫描Entity类(.clsss文件)，使用字节码增强修改这些class文件。
     * @param pkgNames
     */
    @Override
    public void enhance(final String... pkgNames) {
        ClassPool pool = new ClassPool(true);//为了被gc回收,不可使用ClassPool.getDefault()
        //bugfix: javassit高版本类路径查找有问题.
        ClassLoader cl = Thread.currentThread().getContextClassLoader();
        LoaderClassPath lcp = new LoaderClassPath(cl);
        pool.appendClassPath(lcp);
        int n = 0;
        if (roots == null || roots.size() == 0) {
            Set<URL> resources = ClassPathScanningUtils.getPathScanning().getClassResourcesURL(pkgNames);
            for (URL url : resources) {
                if (url.getProtocol().contains(SystemUtils.URL_PROTOCOL_FILE)) {//是否是文件url
                    try {
                        if (enhanceJavassist(url, pool)) {
                            n++;
                        }
                    } catch (Exception e) {
                        logger.error("Enhance error: {}", url, e);
                    }
                }
            }
        } else {
            for (URL root : roots) {
                //设置类加载路径
                ClassPath cp = null;
                try {
                    cp = pool.appendClassPath(root.getFile());
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
                Set<URL> resources = ClassPathScanningUtils.getPathScanning().getClassResourcesURL(root, pkgNames);
                for (URL url : resources) {
                    if (url.getProtocol().contains(SystemUtils.URL_PROTOCOL_FILE)) {//是否是文件url
                        try {
                            if (enhanceJavassist(url, pool)) {
                                n++;
                            }
                        } catch (Exception e) {
                            logger.error("Enhance error: {}", url, e);
                        }
                    }
                }
                //回收路径资源.
                if (cp != null) {
                    pool.removeClassPath(cp);
                }
            }
        }
        pool.removeClassPath(lcp);
        JavassistUtils.removeAllClass(pool);//清理多余的类
        System.out.println(n + " classes enhanced.");
    }

    private boolean enhanceJavassist(URL resource, ClassPool pool) {
        EnhanceTaskJavassist enhancer = new EnhanceTaskJavassist();
        InputStream inputStream = null;
        try {
            inputStream = resource.openStream();
            CtClass cc = pool.makeClass(inputStream);
            byte[] result = enhancer.doEnhance(pool, cc);
            if (result != null) {
                if (result.length == 0) {
                    System.out.println(cc.getName() + " is already enhanced.");
                } else {
                    IOUtils.saveAsFile(new File(resource.toURI()), result);
                    System.out.println("enhanced class:" + cc.getName());// 增强完成
                    return true;
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            IOUtils.closeQuietly(inputStream);
        }
        return false;
    }

}
